import os
import sys

import torch
import torch.utils.data as data

import numpy as np
from PIL import Image
import glob
import random
import cv2

random.seed(1143)


def populate_train_list(low_light_path,paired_normal_path):

    image_list_lowlight = glob.glob(low_light_path + "*.png")
    paired_list = []
    for _, item in enumerate(image_list_lowlight):
        paired_list.append((item, paired_normal_path + item.split('/')[-1]))
    random.shuffle(paired_list)

    return paired_list



class lowlight_loader(data.Dataset):

    def __init__(self, low_light_path, paired_normal_path):

        self.train_list = populate_train_list(low_light_path, paired_normal_path)
        self.size = 256

        self.data_list = self.train_list
        print("Total training examples:", len(self.train_list))




    def __getitem__(self, index):
        data_lowlight_path, data_normal_path = self.data_list[index]
        data_lowlight = Image.open(data_lowlight_path)
        data_lowlight = data_lowlight.resize((self.size,self.size), Image.ANTIALIAS)
        data_lowlight = (np.asarray(data_lowlight)/255.0)
        data_lowlight = torch.from_numpy(data_lowlight).float()
        data_normal = Image.open(data_normal_path)
        data_normal = data_normal.resize((self.size,self.size), Image.ANTIALIAS)
        data_normal = (np.asarray(data_normal)/255.0)
        data_normal = torch.from_numpy(data_normal).float()

        return data_lowlight.permute(2,0,1), data_normal.permute(2,0,1)

    def __len__(self):
        return len(self.data_list)

